# @Time   : 2023/3/29 11:48
# @Author : Fioman
# @Phone  : 13149920693
# @Tips   : Talk is Cheap,Show me the code! ^_^^_^

import datetime
import os

import imutils

import cv2 as cv
import numpy as np
import time

boardAreaFilter = 10000
colorBlue = (255, 0, 0)
colorRed = (0, 0, 255)
colorGreen = (0, 255, 0)
colorList = [colorBlue, colorRed, colorGreen]


def image_show(name, image, showState=False):
    """
    :param name:
    :param image:
    :param showState:
    :return:
    """
    if showState:
        sizeScale = 3
        newWidth = int(image.shape[1] / sizeScale)
        newHeight = int(image.shape[0] / sizeScale)
        showImage = cv.resize(image, (newWidth, newHeight), cv.INTER_AREA)
        cv.namedWindow("{}".format(name))
        cv.imshow("{}".format(name), showImage)
        cv.waitKey(0)


def get_center_by_contour(cnt):
    """
    根据传入的轮廓获取到轮廓的中心,返回轮廓中心的int类型,轮廓中心算出来可能是小数,这里转换为int.
    :param cnt:传入的轮廓
    :return:因为是轮廓的中心,默认返回int类型
    """
    (centerX, centerY), *_ = cv.minAreaRect(cnt)
    return int(centerX + 0.5), int(centerY + 0.5)  # 四舍五入法将浮点数转换为int类型


def get_board_without_angle(image, angleRotated, yMax, extend=300, xSliceStart=0.2, xSliceEnd=0.2):
    """
    获取无角度的物料区域,是否延长,表示是否延长物料区域的四周,保证可以完整的看到物料区域.extend表示延长多少的意思.
    :param image:
    :param boxMaster:
    :param extend:延长多少的标识,0表示不延长
    :return:
    """
    widthHalf = int(image.shape[1] / 2)
    heightHalf = int(image.shape[0] / 2)
    # yStart = max(int(centerY - heightHalf), 0)  # 之所以有零的情况,是因为旋转之后物料框的y起点可能跑到外面去了.
    # yEnd = min(int(centerY + heightHalf), image.shape[0])  # yEnd也是一样的,有可能出去了
    # xStart = max(int(centerX - widthHalf), 0)
    # xEnd = min(int(centerX + widthHalf), image.shape[1])
    M = cv.getRotationMatrix2D((widthHalf, heightHalf), angleRotated, 1)
    yStart = yMax - extend
    yEnd = yMax + 1
    xStart = int(image.shape[1] * xSliceStart)
    xEnd = int(image.shape[1] * xSliceEnd)
    totalBoard = cv.warpAffine(image, M, (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    return totalBoard


def get_two_point_distance(p1, p2):
    """
    获取两点之间的距离,保留小数点后两位
    :param p1: 点1
    :param p2: 点2
    :return:
    """
    assert len(p1) == 2
    assert len(p2) == 2
    ds = np.sqrt((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2)
    return round(ds, 2)


def get_angle_by_k(k):
    """
    根据k值去获取角度,单位是°
    :param k:
    :return:
    """
    if k == 0:
        return 0
    if k is None:
        return 90
    return round(np.arctan(k) * 180 / np.pi, 2)


def get_line_k_and_b(p1, p2):
    """
    根据两个点获取这条线的斜率和斜距
    :param p1: 线段的点1
    :param p2: 线段的点2
    :return:
    """
    if p1[0] == p2[0]:
        # 垂直于x轴,没有k值,斜率是None
        return None, p1[0]
    else:
        # k = (y2 - y1)/(x2-x1)
        # b = y - x * x
        k = (p2[1] - p1[1]) / (p2[0] - p1[0])
        b = p1[1] - k * p1[0]
        return k, b


def get_rect_points_clockwise(rectPoints):
    """
    根据矩形的四个顶点,获取其顺时针方向的四个顶点
    :param rectPoints:
    :return:
    """
    rectPointsSorted = sorted(rectPoints, key=np.sum)
    leftTop = rectPointsSorted[0]
    rightBottom = rectPointsSorted[-1]
    rightTopX = int(rectPointsSorted[1][0])
    leftBottomX = int(rectPointsSorted[2][0])
    if rightTopX > leftBottomX:
        rightTop = rectPointsSorted[1]
        leftBottom = rectPointsSorted[2]
    else:
        rightTop = rectPointsSorted[2]
        leftBottom = rectPointsSorted[1]

    return leftTop, rightTop, rightBottom, leftBottom


def find_ocr_roi(imageOriginal):
    imageGray = cv.cvtColor(imageOriginal, cv.COLOR_BGR2GRAY)
    imageGray[-2:, :] = 0
    image_show("ImageGray", imageGray)
    xStart, xEnd = 1340, 1480
    imageUsed = imageGray[:, xStart:xEnd]
    image_show("ImageUsed",imageUsed)
    thresVal01 = np.mean(imageUsed)
    print("imageUsedThres: {}".format(thresVal01))
    if thresVal01 <= 40:
        thresUsed = 40
    else:
        thresUsed = 50
    T, boardThres = cv.threshold(imageUsed, thresUsed, 255, cv.THRESH_BINARY_INV)
    image_show("Thres", boardThres)

    scharrX = cv.Scharr(boardThres, cv.CV_64F, 1, 0)
    scharrY = cv.Scharr(boardThres, cv.CV_64F, 0, 1)
    scharrX = cv.convertScaleAbs(scharrX)
    scharrY = cv.convertScaleAbs(scharrY)
    scharrXY = cv.addWeighted(scharrX, 0.5, scharrY, 0.5, 0)
    board = scharrXY.copy()
    image_show("ScharrXYEdge", board)
    minLineLenghtEdge = 15
    maxLineGap = 30
    yMin = 100
    yMax = board.shape[0]
    lines = cv.HoughLinesP(board, 1, np.pi / 180, 60, minLineLenghtEdge, maxLineGap)
    boardShow = imageOriginal.copy()
    linesList = []
    colorIndex = 0
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            lineLength = get_two_point_distance((x1, y1), (x2, y2))
            if abs(y2 - y1) < 20 * 3 and abs(x2 - x1) > 5 * 3 and yMin < y1 < yMax and y2 < yMax and lineLength > minLineLenghtEdge:
                x1 += xStart
                x2 += xStart
                linesList.append([x1, y1, (y1 + y2) / 2, lineLength])
                linesList.append([x2, y2, (y1 + y2) / 2, lineLength])
                cv.line(boardShow, (x1, y1), (x2, y2), colorList[colorIndex % 3], 3)
                colorIndex += 1
    linesSortedY = sorted(linesList, key=lambda x: x[2],reverse=True)
    print("LinesFindedSortedY: {}".format(linesSortedY))
    roiHeight = 210
    xSliceStart = 0
    xSliceEnd = 1
    if len(linesSortedY) >= 2:
        lineFinded = linesSortedY[:2]
        p1 = lineFinded[0][0], lineFinded[0][1]
        p2 = lineFinded[1][0], lineFinded[1][1]
        image_show("EdgeFinded", boardShow)
        boardShow = imageOriginal.copy()
        cv.line(boardShow, p1, p2, colorGreen, 3)
        angleFinded = get_angle_by_k(get_line_k_and_b(p1, p2)[0])
        yMax = int(max(p1[1], p2[1]))
        image_show("FinalLineFinded_{}".format(angleFinded), boardShow)
        if yMax <= roiHeight:
            return "fail",imageOriginal
        ocrRoi = get_board_without_angle(imageOriginal, angleFinded, yMax, roiHeight, xSliceStart=xSliceStart, xSliceEnd=xSliceEnd)
        sizeScale = 1
        ocrRoi = cv.resize(ocrRoi, (0, 0), fx=sizeScale, fy=sizeScale, interpolation=cv.INTER_AREA)
        image_show("OCR_ROI", ocrRoi)
        return "ok", ocrRoi
    else:
        # 如果没有找到,就使用另外的方式去找
        cnts = cv.findContours(boardThres, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
        cnts = imutils.grab_contours(cnts)
        cntFinded = sorted(cnts, key=cv.contourArea)
        if len(cntFinded) >= 1:
            cntFinded = cntFinded[-1]
            rectPoints = np.int0(cv.boxPoints(cv.minAreaRect(cntFinded)))
            leftTop, rightTop, rightBottom, leftBottom = get_rect_points_clockwise(rectPoints)
            p1 = leftTop[0] + xStart ,leftTop[1]
            p2 = rightTop[0] + xStart,rightTop[1]
            contourShow = cv.cvtColor(boardThres,cv.COLOR_GRAY2BGR)
            cv.drawContours(contourShow,[rectPoints],-1,colorGreen,3)
            image_show("ContourFinded", contourShow)
            boardShow = imageOriginal.copy()
            cv.line(boardShow, tuple(p1), tuple(p2), colorGreen, 3)
            angleFinded = get_angle_by_k(get_line_k_and_b(p1, p2)[0])
            yMax = int(max(p1[1], p2[1]))
            if yMax <= roiHeight:
                return "fail",imageOriginal
            image_show("FinalLineFinded_{}".format(angleFinded), boardShow)
            ocrRoi = get_board_without_angle(imageOriginal, angleFinded, yMax, roiHeight, xSliceStart=xSliceStart, xSliceEnd=xSliceEnd)
            sizeScale = 1
            ocrRoi = cv.resize(ocrRoi, (0, 0), fx=sizeScale, fy=sizeScale, interpolation=cv.INTER_AREA)
            image_show("OCR_ROI", ocrRoi)
            return "ok", ocrRoi
        else:
            return "fail", imageOriginal


if __name__ == '__main__':
    filePath = r"D:\2022\dongyang_ocr"
    keepDir = filePath
    if not os.path.isdir(keepDir):
        keepDir, _ = os.path.split(filePath)

    keepOkPath = os.path.join(keepDir, "ok")
    keepOkResPath = os.path.join(keepDir, "roi_res")
    keepFailPath = os.path.join(keepDir, "fail")

    dirList = [keepOkPath, keepOkResPath, keepFailPath]
    for dirDemo in dirList:
        if not os.path.exists(dirDemo):
            os.mkdir(dirDemo)

    fileNames = []
    if os.path.isfile(filePath):
        filePath, fileName = os.path.split(filePath)
    else:
        for root, dirs, files in os.walk(filePath):
            if root != filePath:
                continue
            for file in files:
                if file[-4:] != ".bmp":
                    continue
                fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        print("*" * 10 + "第 {} 张图 ({})".format(index + 1, fileName) + "*" * 10)
        calcStart = time.time()
        filePathReal = os.path.join(filePath, fileName)
        imageOriginal = cv.imread(filePathReal, cv.IMREAD_COLOR)
        res, imgRes = find_ocr_roi(imageOriginal)

        if res == "ok":
            print("识别成功!")
            os.remove(filePathReal)
            cv.imwrite(os.path.join(keepOkPath, fileName), imageOriginal)
            cv.imwrite(os.path.join(keepOkResPath, fileName), imgRes)
            calcEnd = time.time()
            print("算法耗时: {:3f}".format(calcEnd - calcStart))
            continue
        else:
            print("识别失败xxxx")
            os.remove(filePathReal)
            cv.imwrite(os.path.join(keepFailPath, fileName), imageOriginal)
            calcEnd = time.time()
            print("算法耗时: {:3f}".format(calcEnd - calcStart))
